import os
import librosa
from fastdtw import fastdtw
from scipy.spatial.distance import euclidean


class AudioRecognizer:
    def __init__(self, template_dir="files/test/", sample_rate=16000, n_mfcc=13):
        self.template_dir = template_dir
        self.sample_rate = sample_rate
        self.n_mfcc = n_mfcc
        self.templates = {}  # 用于存放模板特征的字典

    def load_audio(self, path):
        """加载音频文件并返回音频时序数据及采样率"""
        audio, sr = librosa.load(path, sr=self.sample_rate)
        return audio, sr

    def extract_mfcc_features(self, audio, sr):
        """提取MFCC特征"""
        mfcc = librosa.feature.mfcc(y=audio, sr=sr, n_mfcc=self.n_mfcc)
        mfcc = mfcc.T  # (time_frames, n_mfcc)
        return mfcc

    def build_templates(self):
        """
        从模板文件夹中读取文件并提取MFCC特征。
        假设文件名格式为 <digit>_<index>.wav，比如 1_1.wav, 2_2.wav
        返回结构：dict, key为数字字符串，value为该数字对应的一组MFCC特征列表
        """
        templates = {}
        for fname in os.listdir(self.template_dir):
            if fname.endswith(".wav"):
                digit = fname.split('_')[0]
                path = os.path.join(self.template_dir, fname)
                audio, sr = self.load_audio(path)
                mfcc_features = self.extract_mfcc_features(audio, sr)
                if digit not in templates:
                    templates[digit] = []
                templates[digit].append(mfcc_features)
        self.templates = templates
        return templates

    def dtw_distance(self, mfcc1, mfcc2):
        """
        使用fastdtw计算两个MFCC序列的DTW距离
        mfcc1, mfcc2: shape (T, n_mfcc)
        返回距离值
        """
        distance, _ = fastdtw(mfcc1, mfcc2, dist=euclidean)
        return distance

    def recognize(self, audio_path):
        """
        给定待识别的音频文件路径和模板特征字典，输出识别结果
        需要先调用build_templates()来构建templates
        """
        if not self.templates:
            raise ValueError("Templates not built. Please call build_templates() first.")

        audio, sr = self.load_audio(audio_path)
        test_mfcc = self.extract_mfcc_features(audio, sr)

        best_digit = None
        best_score = float('inf')

        # 遍历每种数字的模板
        for digit, mfcc_list in self.templates.items():
            # 将输入和该类别所有模板比较，取最小距离
            for tmpl_mfcc in mfcc_list:
                dist = self.dtw_distance(test_mfcc, tmpl_mfcc)
                if dist < best_score:
                    best_score = dist
                    best_digit = digit

        return best_digit, best_score


if __name__ == "__main__":
    # 创建识别器实例
    recognizer = AudioRecognizer(template_dir="files/test/")

    # 1. 构建模板库
    recognizer.build_templates()
    print(recognizer.templates)
    # 2. 对待识别音频进行识别
    test_audio_path = "files/test/1.wav"
    digit, score = recognizer.recognize(test_audio_path)
    print("识别结果:", digit, "DTW距离:", score)
